from typing import Union, Dict, Any, Callable
import functools
import json
import math
import numpy as np
from mcp.server.fastmcp import FastMCP

def with_steps(operation: str):
    """
    通用装饰器，为运算函数添加步骤记录功能
    :param operation: 运算符号，如"+", "-"等
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            result = func(*args, **kwargs)
            steps = [f"{args[0]} {operation} {args[1]} = {result}"] if len(args) > 1 else [f"√{args[0]} = {result}"]
            return {"result": result, "steps": steps}
        return wrapper
    return decorator

# 创建MCP服务实例
mcp = FastMCP("Math Calculator")

def format_step(operation: str, a: Any, b: Any, result: Any) -> str:
    """格式化计算步骤"""
    return f"{a} {operation} {b} = {result}"

@mcp.tool()
@with_steps("+")
def add(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """加法运算"""
    return a + b

@mcp.tool()
@with_steps("-")
def subtract(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """减法运算"""
    return a - b

@mcp.tool()
@with_steps("*")
def multiply(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """乘法运算"""
    return a * b

@mcp.tool()
@with_steps("/")
def divide(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """除法运算"""
    if b == 0:
        raise ValueError("除数不能为零")
    return a / b

@mcp.tool()
@with_steps("**")
def power(base: Union[int, float], exponent: Union[int, float]) -> Union[int, float]:
    """幂运算"""
    return base ** exponent

@mcp.tool()
@with_steps("√")
def square_root(x: Union[int, float]) -> float:
    """平方根运算"""
    if x < 0:
        raise ValueError("负数没有实数平方根")
    return x ** 0.5

# Matrix operations
from matrix_operations import (
    matrix_add,
    matrix_multiply,
    matrix_transpose,
    matrix_determinant,
    matrix_inverse,
    matrix_eigen,
    characteristic_polynomial
)

@mcp.tool()
def matrix_add_wrapper(a, b):
    """Add two matrices element-wise"""
    return matrix_add(a, b)

@mcp.tool()
def matrix_multiply_wrapper(a, b):
    """Multiply two matrices"""
    return matrix_multiply(a, b)

@mcp.tool()
def matrix_transpose_wrapper(a):
    """Transpose a matrix"""
    return matrix_transpose(a)

@mcp.tool()
def matrix_determinant_wrapper(a):
    """Calculate determinant of a square matrix"""
    return matrix_determinant(a)

@mcp.tool()
def matrix_inverse_wrapper(a):
    """Calculate inverse of a square matrix"""
    return matrix_inverse(a)

@mcp.tool()
def matrix_eigen_wrapper(a):
    """计算矩阵的特征值和特征向量"""
    return matrix_eigen(a)

@mcp.tool()
def characteristic_polynomial_wrapper(a):
    """计算矩阵的特征多项式系数"""
    return characteristic_polynomial(a)

@mcp.tool()
def minimal_polynomial_wrapper(a):
    """计算矩阵的极小多项式系数"""
    from matrix_operations import minimal_polynomial
    return minimal_polynomial(a)

@mcp.tool()
def matrix_diagonalization_wrapper(a):
    """
    判断矩阵是否可对角化，并返回对角矩阵和变换矩阵
    :param a: 输入矩阵
    :return: 字典包含:
        - diagonalizable: 是否可对角化
        - diagonal_matrix: 对角矩阵(如果可对角化)
        - transformation_matrix: 变换矩阵P(如果可对角化)
    """
    try:
        a = np.array(a)
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square")
            
        # 计算特征值和特征向量
        eigen_result = matrix_eigen(a)
        eigenvalues = np.array(eigen_result["eigenvalues"])
        eigenvectors = np.array(eigen_result["eigenvectors"])
        
        # 检查特征向量矩阵是否可逆(即是否线性无关)
        if np.linalg.matrix_rank(eigenvectors) != a.shape[0]:
            return {"diagonalizable": False}
            
        # 计算对角矩阵和变换矩阵
        diagonal_matrix = np.diag(eigenvalues).tolist()
        transformation_matrix = eigenvectors.tolist()
        
        return {
            "diagonalizable": True,
            "diagonal_matrix": diagonal_matrix,
            "transformation_matrix": transformation_matrix
        }
    except Exception as e:
        raise ValueError(f"Matrix diagonalization error: {str(e)}")

# Basic math operations
def factorial(n: int) -> int:
    """计算阶乘 n!"""
    if n < 0:
        raise ValueError("阶乘只支持非负整数")
    return math.factorial(n)

@mcp.tool()
def combination(n: int, k: int) -> int:
    """计算组合数 C(n,k)"""
    if n < 0 or k < 0:
        raise ValueError("参数必须为非负整数")
    if k > n:
        return 0
    return math.comb(n, k)

@mcp.tool()
def permutation(n: int, k: int) -> int:
    """计算排列数 P(n,k)"""
    if n < 0 or k < 0:
        raise ValueError("参数必须为非负整数")
    if k > n:
        return 0
    return math.perm(n, k)

if __name__ == "__main__":
    mcp.run(transport="stdio")